# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import collections

from absl import app, flags

from coin.base.timestamp import get_timestamp
from coin.strategy.accounting.pnl_balance_printer.factory import get_pnl_balance_printer_cls
from coin.experimental.xguo.fruit.base.strategy_base import (
    BaseStrategy,
    BaseStrategyConfig,
)
from coin.experimental.xguo.fruit.util.fee import get_fee_rate
from coin.experimental.xguo.fruit.util.util import FillStats
# from coin.experimental.xguo.fruit.smm.smm_sim import SmmSimStrategy
from coin.experimental.xguo.fruit.smm.smm_pricer3 import SmmPricer, SmmPricerConfig
from coin.strategy.mm.subscription import SubscriptionRequest
from coin.base.book.util import (
    convert_book_to_dict,
)

FLAGS = flags.FLAGS


OrderFill = collections.namedtuple(
    'OrderFill', ['timestamp', 'fill_price', 'fill_qty', 'product', 'order'])


class SmmStrategy(BaseStrategy):
  def __init__(self, config):
    super().__init__(config)
    self._sim_strat = None  # SmmSimStrategy(config)
    self._pricers = collections.defaultdict(dict)
    sub_req = self.strategy_config.strategy_core_config['trade_subreq']
    self._trade_mea = SubscriptionRequest.from_str(sub_req)

  def prepare_pricers(self):
    for product in self.product_list:
      base = product.base.currency
      if self._trade_mea.exchange == product.exchange:
        self._pricers[base]['trade_symbol'] = product
      else:
        self._pricers[base]['ref_symbol'] = product

    for base, pricer_info in self._pricers.items():
      trade_symbol = pricer_info['trade_symbol']
      ref_symbol = pricer_info.get('ref_symbol')  # can be None
      executor = self.get_executor(trade_symbol)
      exe_config = self.strategy_config.executor_config[trade_symbol]['proto_config']
      config = SmmPricerConfig()
      config.trade_symbol = trade_symbol
      config.ref_symbol = ref_symbol
      config.get_pos = executor.og.get_balance
      config.max_pos = exe_config.max_pos
      config.min_pos = exe_config.min_pos

      pricer = SmmPricer(config)
      self._pricers[base]['pricer'] = pricer
      self._logger.info('Add smm pricer for %s', base)

    # done
    self._logger.info('All pricers prepared! %s', self._pricers)

  def get_pricer(self, product):
    pricer = self._pricers[product.base.currency].get('pricer')
    return pricer

  def install_periodic_task(self):
    self._scheduler.add_periodic_task(
        period=900,
        job_func=self.write_log,
    )
    self._scheduler.add_periodic_task(
        period=30,
        job_func=self.dump_pnl_balance,
    )
    self._scheduler.start()

  def run_before_start(self):
    self.install_periodic_task()
    self.register_balance_printer()

    self.prepare_pricers()
    for pricer_info in self._pricers.values():
      self._logger.info('update tick for korean exchanges!')
      pricer_info['pricer'].update_tick()

    for og in self.order_gateway_list:
      self._logger.info('Update tick formatter for korean exchanges!')
      og.renew_order_formatter()

    for executor in self.executor_list:
      self._logger.info('Update tick for executor!')
      executor.update_tick()

  def write_log(self):
    if not self.is_initialized():
      return
    self._strategy_logger.write_request()

    for og in self.order_gateway_list:
      og.og_logger.write_account_request()

    reserve_proto = self.strategy_config.strategy_core_config['reserve_proto']
    for proto in reserve_proto.values():
      self._strategy_logger.write_reserve(proto)

    for config in self.strategy_config.executor_config.values():
      self._strategy_logger.write_executor_config(config['proto_config'])

  def register_balance_printer(self):
    if self._dry_run:
      return

    prod_list = []
    for prod in self.product_list:
      if prod.exchange == self._trade_mea.exchange:
        prod_list.append(prod)

    core_config = self.strategy_config.strategy_core_config
    reserve_map = list(core_config['reserve'].values())[0]
    og = self.order_gateway_list[0]
    subreq = self.strategy_config.strategy_core_config['subreq_list'][0]
    printer_cls = get_pnl_balance_printer_cls(subreq)
    self._balance_printer = printer_cls(prod_list, self._book_map, reserve_map, logger=self._logger)
    self._balance_printer.register_order_gateway(og)

  def dump_pnl_balance(self):
    if not self.is_initialized():
      return

    for pricer_info in self._pricers.values():
      if not pricer_info['pricer'].is_ready():
        return

    if self._balance_printer is not None:
      self._balance_printer.print_pnl_balance()
      pnl_balances = self._balance_printer.calculate_pnl_balance()
      if pnl_balances is not None:
        for pnl_balance in pnl_balances:
          self._strategy_logger.write_pnl_balance(pnl_balance=pnl_balance,
                                                  quote=None,
                                                  subreq_id=None)

  def on_book(self, product, book):
    if self._sim_strat:
      self._sim_strat.common_on_book(product, book)

    if not self.is_initialized():
      return

    pricer = self.get_pricer(product)
    if pricer is None:
      self._logger.error("%s has no pricer!!!", product)
    pricer.update(product, book)
    if not pricer.is_ready():
      return

    send_order = False
    if pricer.ref_symbol is None:
      send_order = True
    elif product == pricer.ref_symbol:
      send_order = True

    if send_order:
      orders = pricer.gen_order()
      if orders is None:
        self._logger.info('book is not valid!')
        return

      executor = self.get_executor(pricer.trade_symbol)
      assert executor is not None
      agg_buy_order = orders.get('agg_buy')
      agg_sell_order = orders.get('agg_sell')
      mm_order = orders.get('mm')
      self._writer.write(orders, book.timestamp)

      if self._dry_run:
        if self._dry_run == 'VERBOSE':
          self.log_json_message(orders, msg="Submit orders: ")
        return

      if agg_buy_order:
        agg_qty = agg_buy_order['qty']
        if agg_qty > 1e-4:
          executor.force_fire_buy_order(agg_buy_order['price'], agg_qty)

      if agg_sell_order:
        agg_qty = agg_sell_order['qty']
        if agg_qty > 1e-4:
          executor.force_fire_sell_order(agg_sell_order['price'], agg_qty)

      if mm_order:
        executor.manage_mm_orders(
              buy_price=mm_order['pass_buy_price'],
              sell_price=mm_order['pass_sell_price'],
              buy_price_pull=mm_order['buy_price_pull'],
              sell_price_pull=mm_order['sell_price_pull'],
              buy_price_push=mm_order['buy_price_push'],
              sell_price_push=mm_order['sell_price_push'])
    self._writer.write(convert_book_to_dict(book, product=product), book.timestamp)

  def on_trade(self, product, trade_event):
    if self._sim_strat:
      self._sim_strat.on_trade(product, trade_event)

  def on_fill(self, fill_price, fill_qty, order):
    pricer = self.get_pricer(order.product)
    pricer.update_fill(fill_price, fill_qty, order)
    fill_dict = {
        'type': 'fill',
        'fill_price': fill_price,
        'fill_qty': fill_qty,
        'order': order._asdict(),
    }
    self._writer.write(fill_dict)


def main(_):
  config = BaseStrategyConfig.from_file(FLAGS.config)
  strategy = SmmStrategy(config)
  strategy.run()


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Strategy config file.')
  app.run(main)
